import lmdb
import pickle
import os
import logging
from typing import Any, Optional, List, Dict
from contextlib import contextmanager
from FirmwareType import FirmwareType
from UpgradeStatus import UpgradeStatus


class LMDBManager:
    """
    LMDB数据库管理器类
    
    该类负责管理LMDB数据库的连接和所有数据库操作，包括主机信息和固件信息的存储。
    使用LMDB作为底层存储引擎，提供了高性能的键值存储功能。
    """

    def __init__(self, db_path: str = "zijin_upgrade_lmdb", map_size: int = 1024 * 1024 * 1024):  # 1GB
        """
        初始化LMDB数据库管理器
        
        :param db_path: 数据库路径，默认为"zijin_upgrade_lmdb"
        :param map_size: 数据库最大大小，默认为1GB
        """
        # 数据库路径
        self.db_path = db_path
        # 数据库最大映射大小（字节）
        self.map_size = map_size
        # LMDB环境对象
        self.env = None
        # 主机信息子数据库
        self.hosts_db = None
        # 固件信息子数据库
        self.firmware_db = None
        # 计数器子数据库（用于生成唯一ID）
        self.counters_db = None
        # 元数据子数据库
        self.metadata_db = None
        # 日志记录器
        self.logger = logging.getLogger(__name__)
        # 初始化数据库
        self._initialize_db()

    def _initialize_db(self):
        """初始化数据库环境"""
        # 如果数据库目录不存在则创建
        if not os.path.exists(self.db_path):
            os.makedirs(self.db_path)

        # 打开LMDB环境
        self.env = lmdb.open(
            self.db_path,
            map_size=self.map_size,          # 数据库最大大小
            subdir=True,                     # 数据库路径是一个目录
            readonly=False,                  # 数据库可读写
            metasync=True,                   # 元数据同步
            sync=True,                       # 数据同步
            map_async=False,                 # 异步映射
            mode=0o644,                      # 文件权限
            create=True,                     # 如果不存在则创建
            readahead=True,                  # 预读取
            writemap=False,                  # 写映射
            meminit=True,                    # 内存初始化
            max_readers=126,                 # 最大读取器数量
            max_dbs=4,                       # 最大子数据库数量
        )

        # 创建子数据库
        self.hosts_db = self.env.open_db(b'hosts')        # 主机信息数据库
        self.firmware_db = self.env.open_db(b'firmware')  # 固件信息数据库
        self.counters_db = self.env.open_db(b'counters')  # 计数器数据库
        self.metadata_db = self.env.open_db(b'metadata')  # 元数据数据库

        # 初始化计数器（在单独的事务中）
        with self.env.begin(write=True) as txn:
            # 初始化主机ID计数器
            if not txn.get(b'host_id_counter', db=self.counters_db):
                txn.put(b'host_id_counter', pickle.dumps(0), db=self.counters_db)
            # 初始化固件ID计数器
            if not txn.get(b'firmware_id_counter', db=self.counters_db):
                txn.put(b'firmware_id_counter', pickle.dumps(0), db=self.counters_db)

        self.logger.info("LMDB database initialized successfully")

    @contextmanager
    def get_transaction(self, write=False):
        """
        获取数据库事务的上下文管理器
        
        :param write: 是否为写事务，默认为False（只读事务）
        """
        # 开始一个事务
        with self.env.begin(write=write) as txn:
            # 提供事务对象给调用者
            yield txn

    def _get_next_id(self, txn, counter_key: bytes) -> int:
        """
        获取下一个唯一ID
        
        :param txn: 数据库事务对象
        :param counter_key: 计数器键名
        :return: 下一个ID值
        """
        # 从计数器数据库中获取当前ID值
        current_id = pickle.loads(txn.get(counter_key, db=self.counters_db))
        # 计算下一个ID
        next_id = current_id + 1
        # 更新计数器数据库中的值
        txn.put(counter_key, pickle.dumps(next_id), db=self.counters_db)
        return next_id

    def add_host(self, host_ip: str, username: str, password: str) -> Dict[str, Any]:
        """
        添加主机信息到数据库
        
        :param host_ip: 主机IP地址
        :param username: 用户名
        :param password: 密码
        :return: 主机信息字典
        :raises ValueError: 当主机IP已存在时抛出异常
        """
        try:
            # 开始写事务
            with self.env.begin(write=True) as txn:
                # 检查主机是否已存在
                existing_host = self.get_host_by_ip(host_ip)
                if existing_host:
                    error_msg = f"Host with IP {host_ip} already exists"
                    self.logger.warning(error_msg)
                    raise ValueError(error_msg)

                # 获取新的主机ID
                host_id = self._get_next_id(txn, b'host_id_counter')

                # 创建主机记录
                host_data = {
                    'id': host_id,                           # 主机唯一ID
                    'host_ip': host_ip,                      # 主机IP地址
                    'username': username,                    # 用户名
                    'password': password,                    # 密码
                    'upgrade_status': UpgradeStatus.NOTSTARTED,  # 升级状态，默认为未开始
                    'upgrade_log': ''                        # 升级日志，初始为空
                }

                # 保存到数据库
                key = host_ip.encode('utf-8')                # 使用IP作为键
                value = pickle.dumps(host_data)              # 序列化主机数据
                txn.put(key, value, db=self.hosts_db)        # 存储到主机数据库

                self.logger.debug(f"Host added successfully: IP={host_ip}, ID={host_id}")
                return host_data
        except Exception as e:
            self.logger.error(f"Failed to add host: IP={host_ip}, error={str(e)}")
            raise

    def get_host_by_ip(self, host_ip: str) -> Optional[Dict[str, Any]]:
        """
        根据IP地址获取主机信息
        
        :param host_ip: 主机IP地址
        :return: 主机信息字典或None（如果未找到）
        """
        try:
            # 开始只读事务
            with self.env.begin() as txn:
                key = host_ip.encode('utf-8')                # 将IP地址编码为字节
                value = txn.get(key, db=self.hosts_db)       # 从主机数据库获取数据
                if value:
                    host_data = pickle.loads(value)          # 反序列化数据
                    self.logger.debug(f"Host retrieved by IP: {host_ip}")
                    return host_data
                self.logger.debug(f"Host not found by IP: {host_ip}")
                return None
        except Exception as e:
            self.logger.error(f"Failed to get host by IP: {host_ip}, error={str(e)}")
            return None

    def get_host_by_id(self, host_id: int) -> Optional[Dict[str, Any]]:
        """
        根据ID获取主机信息
        
        :param host_id: 主机ID
        :return: 主机信息字典或None（如果未找到）
        """
        try:
            # 开始只读事务
            with self.env.begin() as txn:
                # 创建游标遍历主机数据库
                cursor = txn.cursor(db=self.hosts_db)
                for key, value in cursor:
                    host_data = pickle.loads(value)          # 反序列化数据
                    if host_data['id'] == host_id:           # 匹配ID
                        self.logger.debug(f"Host retrieved by ID: {host_id}")
                        return host_data
                self.logger.debug(f"Host not found by ID: {host_id}")
                return None
        except Exception as e:
            self.logger.error(f"Failed to get host by ID: {host_id}, error={str(e)}")
            return None

    def get_all_hosts(self) -> List[Dict[str, Any]]:
        """
        获取所有主机信息
        
        :return: 主机信息列表，按ID排序
        """
        try:
            hosts = []
            # 开始只读事务
            with self.env.begin() as txn:
                # 创建游标遍历主机数据库
                cursor = txn.cursor(db=self.hosts_db)
                for key, value in cursor:
                    host_data = pickle.loads(value)          # 反序列化数据
                    hosts.append(host_data)
            # 按ID排序
            hosts.sort(key=lambda x: x['id'])
            self.logger.debug(f"Retrieved all hosts, count: {len(hosts)}")
            return hosts
        except Exception as e:
            self.logger.error(f"Failed to get all hosts, error={str(e)}")
            return []

    def update_host(self, host_id: int, host_ip: str, username: str, password: str) -> bool:
        """
        更新主机信息
        
        :param host_id: 主机ID
        :param host_ip: 新的主机IP地址
        :param username: 新的用户名
        :param password: 新的密码
        :return: 是否更新成功
        """
        try:
            # 开始写事务
            with self.env.begin(write=True) as txn:
                # 查找主机
                cursor = txn.cursor(db=self.hosts_db)
                host_data = None
                host_key = None

                # 遍历查找匹配ID的主机
                for key, value in cursor:
                    data = pickle.loads(value)
                    if data['id'] == host_id:
                        host_data = data
                        host_key = key
                        break

                # 如果未找到主机，返回False
                if not host_data:
                    self.logger.warning(f"Host not found for update: ID={host_id}")
                    return False

                # 更新主机信息
                old_ip = host_data['host_ip']
                host_data['host_ip'] = host_ip
                host_data['username'] = username
                host_data['password'] = password

                # 保存更新后的数据
                value = pickle.dumps(host_data)
                txn.put(host_key, value, db=self.hosts_db)

                self.logger.debug(f"Host updated successfully: ID={host_id}, Old IP={old_ip}, New IP={host_ip}")
                return True
        except Exception as e:
            self.logger.error(f"Failed to update host: ID={host_id}, error={str(e)}")
            return False

    def update_host_status(self, host_ip: str, status: UpgradeStatus, log: str = '') -> bool:
        """
        更新主机升级状态和日志
        
        :param host_ip: 主机IP地址
        :param status: 新的升级状态
        :param log: 升级日志
        :return: 是否更新成功
        """
        try:
            # 开始写事务
            with self.env.begin(write=True) as txn:
                key = host_ip.encode('utf-8')
                value = txn.get(key, db=self.hosts_db)
                # 如果未找到主机，返回False
                if not value:
                    self.logger.warning(f"Host not found for status update: IP={host_ip}")
                    return False

                # 反序列化主机数据
                host_data = pickle.loads(value)
                old_status = host_data.get('upgrade_status', 'Unknown')
                # 更新状态和日志
                host_data['upgrade_status'] = status
                host_data['upgrade_log'] = log

                # 序列化并保存更新后的数据
                updated_value = pickle.dumps(host_data)
                txn.put(key, updated_value, db=self.hosts_db)

                self.logger.debug(
                    f"Host status updated: IP={host_ip}, Old Status={old_status}, New Status={status.name}, Log Length={len(log)}")
                return True
        except Exception as e:
            self.logger.error(f"Failed to update host status: IP={host_ip}, error={str(e)}")
            return False

    def delete_host(self, host_id: int) -> bool:
        """
        删除主机信息
        
        :param host_id: 主机ID
        :return: 是否删除成功
        """
        try:
            # 开始写事务
            with self.env.begin(write=True) as txn:
                # 创建游标遍历主机数据库
                cursor = txn.cursor(db=self.hosts_db)
                for key, value in cursor:
                    host_data = pickle.loads(value)
                    # 匹配主机ID
                    if host_data['id'] == host_id:
                        ip = host_data['host_ip']
                        # 删除主机记录
                        result = txn.delete(key, db=self.hosts_db)
                        if result:
                            self.logger.debug(f"Host deleted successfully: ID={host_id}, IP={ip}")
                        else:
                            self.logger.warning(f"Failed to delete host: ID={host_id}, IP={ip}")
                        return result
                self.logger.warning(f"Host not found for deletion: ID={host_id}")
                return False
        except Exception as e:
            self.logger.error(f"Failed to delete host: ID={host_id}, error={str(e)}")
            return False

    def add_firmware(self, file_name: str, file_path: str, file_size: int, firmware_type: FirmwareType) -> Dict[
        str, Any]:
        """
        添加固件信息到数据库
        
        :param file_name: 固件文件名
        :param file_path: 固件文件路径
        :param file_size: 固件文件大小
        :param firmware_type: 固件类型
        :return: 固件信息字典
        :raises ValueError: 当固件文件名已存在时抛出异常
        """
        try:
            # 开始写事务
            with self.env.begin(write=True) as txn:
                # 检查固件是否已存在
                existing_firmware = self.get_firmware_by_name(file_name)
                if existing_firmware:
                    error_msg = f"Firmware with name {file_name} already exists"
                    self.logger.warning(error_msg)
                    raise ValueError(error_msg)

                # 获取新的固件ID
                firmware_id = self._get_next_id(txn, b'firmware_id_counter')

                # 创建固件记录
                import datetime
                firmware_data = {
                    'id': firmware_id,                       # 固件唯一ID
                    'file_name': file_name,                  # 固件文件名
                    'file_path': file_path,                  # 固件文件路径
                    'file_size': file_size,                  # 固件文件大小
                    'upload_time': datetime.datetime.now(),  # 上传时间
                    'type': firmware_type                    # 固件类型
                }

                # 保存到数据库
                key = file_name.encode('utf-8')              # 使用文件名作为键
                value = pickle.dumps(firmware_data)          # 序列化固件数据
                txn.put(key, value, db=self.firmware_db)     # 存储到固件数据库

                self.logger.debug(f"Firmware added successfully: Name={file_name}, ID={firmware_id}, Size={file_size}")
                return firmware_data
        except Exception as e:
            self.logger.error(f"Failed to add firmware: Name={file_name}, error={str(e)}")
            raise

    def get_firmware_by_name(self, file_name: str) -> Optional[Dict[str, Any]]:
        """
        根据文件名获取固件信息
        
        :param file_name: 固件文件名
        :return: 固件信息字典或None（如果未找到）
        """
        try:
            # 开始只读事务
            with self.env.begin() as txn:
                key = file_name.encode('utf-8')              # 将文件名编码为字节
                value = txn.get(key, db=self.firmware_db)    # 从固件数据库获取数据
                if value:
                    firmware_data = pickle.loads(value)      # 反序列化数据
                    self.logger.debug(f"Firmware retrieved by name: {file_name}")
                    return firmware_data
                self.logger.debug(f"Firmware not found by name: {file_name}")
                return None
        except Exception as e:
            self.logger.error(f"Failed to get firmware by name: {file_name}, error={str(e)}")
            return None

    def get_firmware_by_id(self, firmware_id: int) -> Optional[Dict[str, Any]]:
        """
        根据ID获取固件信息
        
        :param firmware_id: 固件ID
        :return: 固件信息字典或None（如果未找到）
        """
        try:
            # 开始只读事务
            with self.env.begin() as txn:
                # 创建游标遍历固件数据库
                cursor = txn.cursor(db=self.firmware_db)
                for key, value in cursor:
                    firmware_data = pickle.loads(value)      # 反序列化数据
                    if firmware_data['id'] == firmware_id:   # 匹配ID
                        self.logger.debug(f"Firmware retrieved by ID: {firmware_id}")
                        return firmware_data
                self.logger.debug(f"Firmware not found by ID: {firmware_id}")
                return None
        except Exception as e:
            self.logger.error(f"Failed to get firmware by ID: {firmware_id}, error={str(e)}")
            return None

    def get_all_firmware(self) -> List[Dict[str, Any]]:
        """
        获取所有固件信息
        
        :return: 固件信息列表
        """
        try:
            firmware_list = []
            # 开始只读事务
            with self.env.begin() as txn:
                # 创建游标遍历固件数据库
                cursor = txn.cursor(db=self.firmware_db)
                for key, value in cursor:
                    firmware_data = pickle.loads(value)      # 反序列化数据
                    firmware_list.append(firmware_data)
            self.logger.debug(f"Retrieved all firmware, count: {len(firmware_list)}")
            return firmware_list
        except Exception as e:
            self.logger.error(f"Failed to get all firmware, error={str(e)}")
            return []

    def delete_firmware(self, firmware_id: int) -> bool:
        """
        删除固件信息
        
        :param firmware_id: 固件ID
        :return: 是否删除成功
        """
        try:
            # 开始写事务
            with self.env.begin(write=True) as txn:
                # 创建游标遍历固件数据库
                cursor = txn.cursor(db=self.firmware_db)
                for key, value in cursor:
                    firmware_data = pickle.loads(value)
                    # 匹配固件ID
                    if firmware_data['id'] == firmware_id:
                        name = firmware_data['file_name']
                        # 删除固件记录
                        result = txn.delete(key, db=self.firmware_db)
                        if result:
                            self.logger.debug(f"Firmware deleted successfully: ID={firmware_id}, Name={name}")
                        else:
                            self.logger.warning(f"Failed to delete firmware: ID={firmware_id}, Name={name}")
                        return result
                self.logger.warning(f"Firmware not found for deletion: ID={firmware_id}")
                return False
        except Exception as e:
            self.logger.error(f"Failed to delete firmware: ID={firmware_id}, error={str(e)}")
            return False

    def close(self):
        """关闭数据库连接"""
        if self.env:
            self.env.close()  # 关闭LMDB环境
            self.logger.debug("LMDB database connection closed")

    def __enter__(self):
        """
        上下文管理器入口
        
        :return: LMDBManager实例
        """
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        上下文管理器出口，确保数据库连接被正确关闭
        
        :param exc_type: 异常类型
        :param exc_val: 异常值
        :param exc_tb: 异常回溯
        """
        self.close()